package com.squareup.anvil.compiler.internal.testing

import com.squareup.anvil.compiler.api.AnvilContext
import com.squareup.anvil.compiler.api.CodeGenerator
import com.squareup.anvil.compiler.api.FileWithContent
import com.squareup.anvil.compiler.api.createGeneratedFile
import com.squareup.anvil.compiler.internal.reference.ClassReference.Psi
import com.squareup.anvil.compiler.internal.reference.classAndInnerClassReferences
import com.squareup.anvil.compiler.internal.testing.SimpleCodeGenerator.CodeGeneratorAction
import com.squareup.anvil.compiler.internal.testing.SimpleSourceFileTrackingBehavior.NO_SOURCE_TRACKING
import com.squareup.anvil.compiler.internal.testing.SimpleSourceFileTrackingBehavior.TRACKING_WITH_NO_SOURCES
import com.squareup.anvil.compiler.internal.testing.SimpleSourceFileTrackingBehavior.TRACK_SOURCE_FILES
import org.jetbrains.kotlin.descriptors.ModuleDescriptor
import org.jetbrains.kotlin.psi.KtFile
import java.io.File

private var counter = 0

public fun simpleCodeGenerator(
  trackSourceFiles: SimpleSourceFileTrackingBehavior = TRACK_SOURCE_FILES,
  applicable: (AnvilContext) -> Boolean = { true },
  mapper: SimpleCodeGenerator.(clazz: Psi) -> String?,
): SimpleCodeGenerator = SimpleCodeGenerator(applicable) { codeGenDir, module, projectFiles ->

  projectFiles
    .classAndInnerClassReferences(module)
    .mapNotNull { clazz ->
      val content = mapper.invoke(this, clazz) ?: return@mapNotNull null
      val (packageName, fileName) = parseSimpleFileContents(content)

      when (trackSourceFiles) {
        NO_SOURCE_TRACKING ->
          @Suppress("DEPRECATION")
          createGeneratedFile(
            codeGenDir = codeGenDir,
            packageName = packageName,
            fileName = fileName,
            content = content,
          )

        TRACKING_WITH_NO_SOURCES,
        TRACK_SOURCE_FILES,
        -> createGeneratedFile(
          codeGenDir = codeGenDir,
          packageName = packageName,
          fileName = fileName,
          content = content,
          sourceFiles = if (trackSourceFiles == TRACKING_WITH_NO_SOURCES) {
            emptySet()
          } else {
            setOf(clazz.containingFileAsJavaFile)
          },
        )
      }
    }
    .toList()
}

public fun simpleCodeGenerator(
  applicable: (AnvilContext) -> Boolean = { true },
  generateCode: CodeGeneratorAction<SimpleCodeGenerator>,
): SimpleCodeGenerator = SimpleCodeGenerator(applicable, generateCode)

public class SimpleCodeGenerator internal constructor(
  private val applicable: (AnvilContext) -> Boolean = { true },
  private val codeGeneratorAction: CodeGeneratorAction<SimpleCodeGenerator>,
) : CodeGenerator {
  public var isApplicableCalls: Int = 0
    private set
  internal val inputFileContentToGenerateCalls: MutableMap<String, Int> = mutableMapOf()

  public fun getGenerateCallsForInputFileContent(content: String): Int =
    inputFileContentToGenerateCalls[content] ?: 0

  override fun isApplicable(context: AnvilContext): Boolean {
    isApplicableCalls++
    return applicable(context)
  }

  override fun generateCode(
    codeGenDir: File,
    module: ModuleDescriptor,
    projectFiles: Collection<KtFile>,
  ): Collection<FileWithContent> {
    for (file in projectFiles) {
      inputFileContentToGenerateCalls[file.text] =
        getGenerateCallsForInputFileContent(file.text) + 1
    }
    return with(codeGeneratorAction) {
      generate(
        codeGenDir = codeGenDir,
        module = module,
        projectFiles = projectFiles,
      )
    }
  }

  public fun interface CodeGeneratorAction<T : CodeGenerator> {
    public fun T.generate(
      codeGenDir: File,
      module: ModuleDescriptor,
      projectFiles: Collection<KtFile>,
    ): Collection<FileWithContent>
  }
}

/**
 * Defines the kind of [FileWithContent] that should be generated by [simpleCodeGenerator].
 */
public enum class SimpleSourceFileTrackingBehavior {
  /**
   * Returns the deprecated [GeneratedFile][com.squareup.anvil.compiler.api.GeneratedFile]
   * that is unsupported when the [trackSourceFiles][AnvilContext.trackSourceFiles]
   * flag is set.
   */
  NO_SOURCE_TRACKING,

  /**
   * Returns a [GeneratedFileWithSources][com.squareup.anvil.compiler.api.GeneratedFileWithSources],
   * but with an empty list of source files.
   */
  TRACKING_WITH_NO_SOURCES,

  /**
   * Returns a [GeneratedFileWithSources][com.squareup.anvil.compiler.api.GeneratedFileWithSources]
   * with the source file of the class that is being processed.
   */
  TRACK_SOURCE_FILES,
}

/**
 * Represents the metadata of a file that is generated by [simpleCodeGenerator].
 */
public data class SimpleFileContents(
  val packageName: String,
  val fileName: String,
)

/**
 * Parses [SimpleFileContents] metadata from a given source [content].
 */
public fun parseSimpleFileContents(content: String): SimpleFileContents {
  val packageName = content.lines()
    .map { it.trim() }
    .firstNotNullOfOrNull { line ->
      line
        .takeIf { it.startsWith("package ") }
        ?.substringAfter("package ")
    }
    ?: "tempPackage"

  val fileName = content.lines()
    .map { it.trim() }
    .firstNotNullOfOrNull { line ->
      // Try finding the class name.
      line
        .takeIf { it.startsWith("class ") || it.contains(" class ") }
        ?.substringAfter("class ")
        ?.trim()
        ?.substringBefore(" ")
        ?: line
          // Check for interfaces, too.
          .takeIf { it.startsWith("interface ") || it.contains(" interface ") }
          ?.substringAfter("interface ")
          ?.trim()
          ?.substringBefore(" ")
    }
    ?: "NewFile${counter++}"

  return SimpleFileContents(
    packageName = packageName,
    fileName = fileName,
  )
}
